# code: utf-8
import copy
import math

import numpy as np

try:
    import DetailDesignV2.jxq.group as group
except:
    import group as group
import itertools as it
from numpy import linalg as LA

try:
    import DetailDesignV2.jxq.model_operator as model_operator
except:
    import model_operator as model_operator

try:
    from Sympgroup import SymGroup, SymBody
except:
    from .Sympgroup import SymGroup, SymBody


# 调整RPS点的顺序，与距离最近的clamp点靠近
def modify_rps_order(obj, rps_coor):
    if len(rps_coor) > 1:
        clamp_points = get_clamp_point(obj)
        clamp_points_list = [v[0] for v in clamp_points]
        rps_combination_dict = []
        rps_dist_list = []
        for temp_rps in it.permutations(rps_coor, len(rps_coor)):
            dis_rps_clamp = 0.0
            for clamp_point, rps_point in zip(clamp_points_list, temp_rps):
                dis_rps_clamp += LA.norm(clamp_point - rps_point)
            rps_dist_list.append(dis_rps_clamp)
            rps_combination_dict.append(temp_rps)
        min_index = rps_dist_list.index(min(rps_dist_list))
        rps_coor = rps_combination_dict[min_index]
    print("调整后的rps坐标为：", rps_coor)
    return rps_coor


# 获取夹具点的坐标
def get_clamp_point(clamp=None):
    point = []
    clamp_point = clamp.obj.get('clamp_points', [])
    for i in range(len(clamp_point)):
        for j in range(len(clamp_point[i])):
            point.append(clamp.get_point(clamp_point[i][j]))
    # print(point)
    points = np.array(point).reshape(len(clamp_point), len(clamp_point[0]), 3)
    # print("夹紧点的坐标为：", points)
    return points


# 替换符号中带关系的部分
def replaceRelationPara(total_loss_sym, possible_para_sym_list, sym_dict):
    new_total_loss_sym = copy.deepcopy(total_loss_sym)
    for e in possible_para_sym_list:
        for k, v in e.items():
            if 'relation' in v['v']:
                exec(v['v']['relation'].strip(), sym_dict)

    for e in possible_para_sym_list:
        for k, v in e.items():
            if 'relation' in v['v']:
                new_total_loss_sym = new_total_loss_sym.subs(v['sym'], sym_dict[k])
    return new_total_loss_sym, sym_dict


# 替换符号中的常量
def replaceConstPara(total_loss_sym, const_para_sym_list):
    new_total_loss_sym = copy.deepcopy(total_loss_sym)
    for e in const_para_sym_list:
        for k, v in e.items():
            new_total_loss_sym = new_total_loss_sym.subs(v['sym'], v['v']['v'])
    return new_total_loss_sym


# 夹具各部件尺寸更新
def update_args(obj, rps_coor):
    import sympy as sym
    from scipy.optimize import minimize

    sym_obj = SymBody(obj.obj)
    sym_obj = model_operator.update_model_by_history(sym_obj, obj)

    clamp_points = get_clamp_point(sym_obj)
    loss_sym = [sym.sqrt(sum((clamp_point - rps_point) * (clamp_point - rps_point).T)) for
                clamp_point_list, rps_point in zip(clamp_points, rps_coor) for clamp_point in clamp_point_list]

    total_loss_sym = sum(loss_sym)

    para_sym_list = [(sym_obj.syms[e], e) for e in sym_obj.args]
    para_dict = copy.deepcopy(sym_obj.args)
    sym_dict = copy.deepcopy(sym_obj.syms)

    possible_para_sym_list = [{e[1]: {"v": sym_obj.args[e[1]], "sym": sym_dict[e[1]]}} for e in para_sym_list
                              if 'isconst' not in sym_obj.args[e[1]] or not sym_obj.args[e[1]]['isconst']]

    total_loss_sym, sym_dict = replaceRelationPara(total_loss_sym, possible_para_sym_list, sym_dict)

    relation_para_sym_list = [{k: {"v": sym_obj.args[k], "sym": sym_dict[k]}} for e in possible_para_sym_list
                              for k, v in e.items() if 'relation' in v['v']]

    need_para_sym_list = [{k: {"v": sym_obj.args[k], "sym": sym_dict[k]}} for e in possible_para_sym_list
                          for k, v in e.items() if 'relation' not in v['v']]

    const_para_sym_list = [{e[1]: {"v": sym_obj.args[e[1]], "sym": sym_dict[e[1]]}} for e in para_sym_list
                           if 'isconst' in para_dict[e[1]] and para_dict[e[1]]['isconst']]

    const_para_val = [{'name': e[1], 'v': sym_obj.args[e[1]]['v']} for e in para_sym_list
                      if 'isconst' in para_dict[e[1]] and para_dict[e[1]]['isconst']]

    if len(const_para_sym_list) > 0:
        total_loss_sym = replaceConstPara(total_loss_sym, const_para_sym_list)

    possible_para_sym = [v['sym'] for e in need_para_sym_list for k, v in e.items()]

    total_loss_fun = sym.lambdify(([possible_para_sym]), total_loss_sym, 'numpy')

    jacobian_fun = sym.Matrix([total_loss_sym]).jacobian(possible_para_sym)
    jacobian_fun_list = [sym.lambdify(([possible_para_sym]), e, 'numpy') for e in jacobian_fun[0, :]]

    def jacob(x):
        result = [e(x) for e in jacobian_fun_list]
        return np.array(result)

    possible_value = [v['v']['v'] for e in need_para_sym_list for k, v in e.items()]
    possible_bounds = [v['v']['range'] if 'range' in v['v'] else [0, None] for e in need_para_sym_list for k, v in
                       e.items()]

    to_solve_problem = [{"v": possible_value, "bounds": possible_bounds}]

    for i in range(len(need_para_sym_list)):
        e = need_para_sym_list[i]
        for k, v in e.items():
            if 'step' in v['v']:
                if v['v']['step'] is not None and isinstance(v['v']['step'], list):
                    increase = []
                    for step in v['v']['step']:
                        for problem in to_solve_problem:
                            problem["bounds"][i] = [step, step]
                            problem["v"][i] = step
                            increase.append(copy.deepcopy(problem))
                    to_solve_problem = increase

    initial_result_list = [
        minimize(total_loss_fun, problem['v'], method='L-BFGS-B', jac=jacob, bounds=problem['bounds']) for problem in
        to_solve_problem]

    new_to_solve_problem = copy.deepcopy(zip(initial_result_list, to_solve_problem))
    for index in range(len(need_para_sym_list)):
        for k, v in need_para_sym_list[index].items():
            if 'step' in v['v']:
                if v['v']['step'] is not None and not isinstance(v['v']['step'], list):
                    increase = []
                    for res, problem in new_to_solve_problem:
                        near_step = math.ceil((res.x[index] - v['v']['range'][0]) / v['v']['step'])
                        for next in range(-1, 2):
                            val = v['v']['range'][0] + (near_step + next) * v['v']['step']
                            problem["v"][index] = val
                            problem['bounds'][index] = [val, val]
                            increase.append([res, copy.deepcopy(problem)])
                    new_to_solve_problem = increase

    finally_result_list = [
        minimize(total_loss_fun, problem['v'], method='L-BFGS-B', jac=jacob, bounds=problem['bounds'])
        for res, problem in new_to_solve_problem]
    res = min(finally_result_list, key=lambda x: x.fun)

    need_para_val = [{'name': k, 'v': res.x[i]} for i in range(len(res.x)) for k, v in need_para_sym_list[i].items()]

    relation_para_val = []

    for relation in relation_para_sym_list:
        for k, v in relation.items():
            for para in need_para_val:
                v['sym'] = v['sym'].subs(sym_dict[para['name']], para['v'])
            relation_para_val.append({'name': k, 'v': float(v['sym'])})
    new_para = const_para_val + need_para_val + relation_para_val

    obj.obj['args'] = new_para
    new_obj = group.Body(obj.obj)
    obj = model_operator.update_model_by_history(new_obj, obj)
    return obj


def get_rps_diff(obj, rps_coor):
    total_loss = 0
    clamp_points = get_clamp_point(obj)
    for clamp_point_list, rps_point in zip(clamp_points, rps_coor):
        new_rps_list = [rps_point for i in range(len(clamp_point_list))]
        for clamp_point, rps_point in zip(clamp_point_list, new_rps_list):
            total_loss += LA.norm(clamp_point - rps_point)
    return total_loss


# 模型更新
def model_update(obj, rps_coor, env=None, section=None, air_point=None, lr=1, dx=1, iters=1):
    # 调整RPS点的顺序，与距离最近的clamp点靠近
    rps_coor = modify_rps_order(obj, rps_coor)
    print("调整后的rps_coor：", rps_coor)

    obj = update_args(obj, rps_coor)
    print("更新后夹具点的坐标为：{}，距离为：{}".format(get_clamp_point(obj), get_rps_diff(obj, rps_coor)))
    print("更新后的零件尺寸为：", obj.args)
    # print("模型更新时间为：", time.time() - start_time)
    return obj


if __name__ == "__main__":
    rps1 = [[-370.5, 0, -95]]
    #
    # rps2 = [
    #         [-370.5, 80, -95],
    #         [-370.5, -80, -95],
    #     ]
    #
    # rps3 = [
    #
    #         [-185.0,  0, -90],
    #         [-360.5, 80, -95],
    #         [-360.5, -80, -95],
    #     ]
    # # Env = trimesh.load("work_2_xy.stl")
    # x = group.Group("clamp")
    # # x.show()
    # section = [1, 0, 0]
    # air_point = [0, 0, 0]
    # model_update(obj=x, rps_coor=rps1, env=None, section=None, air_point=None, lr=1, dx=0.5, iters=100)
